#!/usr/bin/env python3
"""
PiDiNet Android 导出便捷脚本
检查环境依赖并执行模型导出
"""

import sys
import subprocess
import importlib
import os
from pathlib import Path

def check_python_version():
    """检查Python版本"""
    print("🔍 检查Python版本...")
    if sys.version_info < (3, 7):
        print("❌ 需要Python 3.7或更高版本")
        print(f"   当前版本: {sys.version}")
        return False
    print(f"✅ Python版本: {sys.version.split()[0]}")
    return True

def check_dependencies():
    """检查必需的依赖"""
    print("🔍 检查依赖包...")
    
    required_packages = {
        'torch': 'PyTorch',
        'torchvision': 'TorchVision', 
        'cv2': 'OpenCV (opencv-python)',
        'numpy': 'NumPy',
        'PIL': 'Pillow',
        'onnx': 'ONNX',
        'onnxruntime': 'ONNX Runtime'
    }
    
    missing_packages = []
    
    for package, name in required_packages.items():
        try:
            importlib.import_module(package)
            print(f"✅ {name}")
        except ImportError:
            print(f"❌ {name} - 未安装")
            missing_packages.append(package)
    
    return missing_packages

def install_dependencies(missing_packages):
    """安装缺失的依赖"""
    if not missing_packages:
        return True
    
    print(f"\n📦 需要安装 {len(missing_packages)} 个依赖包...")
    
    # 映射包名到pip安装名
    pip_names = {
        'torch': 'torch',
        'torchvision': 'torchvision', 
        'cv2': 'opencv-python',
        'numpy': 'numpy',
        'PIL': 'Pillow',
        'onnx': 'onnx',
        'onnxruntime': 'onnxruntime'
    }
    
    try:
        for package in missing_packages:
            pip_name = pip_names.get(package, package)
            print(f"正在安装 {pip_name}...")
            
            result = subprocess.run([
                sys.executable, '-m', 'pip', 'install', pip_name
            ], capture_output=True, text=True)
            
            if result.returncode == 0:
                print(f"✅ {pip_name} 安装成功")
            else:
                print(f"❌ {pip_name} 安装失败:")
                print(result.stderr)
                return False
        
        return True
        
    except Exception as e:
        print(f"❌ 安装过程中出错: {e}")
        return False

def check_model_file():
    """检查模型文件是否存在"""
    print("🔍 检查模型文件...")
    model_path = "table5_pidinet-tiny-l.pth"
    
    if os.path.exists(model_path):
        file_size = os.path.getsize(model_path) / (1024 * 1024)  # MB
        print(f"✅ 模型文件存在: {model_path} ({file_size:.1f} MB)")
        return True
    else:
        print(f"❌ 模型文件不存在: {model_path}")
        print("   请确保模型文件在当前目录中")
        return False

def run_export():
    """运行模型导出"""
    print("\n🚀 开始导出Android模型...")
    
    try:
        from android_export import AndroidModelExporter
        
        exporter = AndroidModelExporter("table5_pidinet-tiny-l.pth")
        exporter.export_all()
        
        return True
        
    except Exception as e:
        print(f"❌ 导出失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def show_next_steps():
    """显示后续步骤"""
    print("\n" + "="*50)
    print("🎉 导出完成！接下来的步骤:")
    print("="*50)
    
    print("\n📁 生成的文件:")
    
    # 检查生成的文件
    files_to_check = [
        ("models/pidinet_mobile.onnx", "ONNX模型文件 (主要)"),
        ("models/model_config.json", "模型配置文件"),
        ("android/DocumentEdgeDetector.java", "Android检测类"),
        ("android/MainActivity.java", "使用示例"),
        ("android/dependencies.gradle", "依赖配置")
    ]
    
    for file_path, description in files_to_check:
        if os.path.exists(file_path):
            print(f"✅ {file_path} - {description}")
        else:
            print(f"❌ {file_path} - 生成失败")
    
    print("\n📋 Android集成步骤:")
    print("1. 将 models/ 目录中的文件复制到 Android 项目的 assets/ 目录")
    print("2. 在 build.gradle 中添加 dependencies.gradle 中的依赖")
    print("3. 复制 DocumentEdgeDetector.java 到你的项目")
    print("4. 参考 MainActivity.java 实现相机集成")
    print("5. 查看 README_Android集成指南.md 获取详细说明")

def main():
    """主函数"""
    print("🔧 PiDiNet Android 导出工具")
    print("="*40)
    
    # 检查Python版本
    if not check_python_version():
        return False
    
    # 检查模型文件
    if not check_model_file():
        return False
    
    # 检查依赖
    missing_packages = check_dependencies()
    
    if missing_packages:
        print(f"\n⚠️  缺少 {len(missing_packages)} 个依赖包")
        user_input = input("是否自动安装缺失的依赖? (y/n): ").lower().strip()
        
        if user_input in ['y', 'yes', '是']:
            if not install_dependencies(missing_packages):
                print("\n❌ 依赖安装失败，请手动安装后重试")
                return False
        else:
            print("\n📋 请手动安装以下依赖:")
            for package in missing_packages:
                pip_names = {
                    'torch': 'torch',
                    'torchvision': 'torchvision', 
                    'cv2': 'opencv-python',
                    'numpy': 'numpy',
                    'PIL': 'Pillow',
                    'onnx': 'onnx',
                    'onnxruntime': 'onnxruntime'
                }
                pip_name = pip_names.get(package, package)
                print(f"   pip install {pip_name}")
            return False
    
    # 运行导出
    if run_export():
        show_next_steps()
        return True
    else:
        print("\n❌ 导出失败")
        return False

if __name__ == "__main__":
    success = main()
    
    if success:
        print("\n🎊 所有任务完成！")
    else:
        print("\n💡 如果遇到问题，请检查:")
        print("   1. Python环境是否正确")
        print("   2. 模型文件是否存在")
        print("   3. 网络连接是否正常 (用于下载依赖)")
        print("   4. 查看错误信息进行调试")
    
    input("\n按回车键退出...") 